/**
 * @file test_arma_base.cpp
 *
 * Unit tests for basic extensions of Armadillo
 *
 * @author Dahua Lin
 */


#include "test_common.h"

using namespace gml;

/************************************************
 *
 *  Auxiliary functions
 *
 ************************************************/

inline mat bsx_add_rows(const mat& X, const rowvec& row)
{
	return X + arma::repmat(row, X.n_rows, 1);
}

inline mat bsx_add_cols(const mat& X, const colvec& col)
{
	return X + arma::repmat(col, 1, X.n_cols);
}

inline mat bsx_sub_rows(const mat& X, const rowvec& row)
{
	return X - arma::repmat(row, X.n_rows, 1);
}

inline mat bsx_sub_cols(const mat& X, const colvec& col)
{
	return X - arma::repmat(col, 1, X.n_cols);
}

inline mat bsx_mul_rows(const mat& X, const rowvec& row)
{
	return X % arma::repmat(row, X.n_rows, 1);
}

inline mat bsx_mul_cols(const mat& X, const colvec& col)
{
	return X % arma::repmat(col, 1, X.n_cols);
}

inline mat bsx_div_rows(const mat& X, const rowvec& row)
{
	return X / arma::repmat(row, X.n_rows, 1);
}

inline mat bsx_div_cols(const mat& X, const colvec& col)
{
	return X / arma::repmat(col, 1, X.n_cols);
}



inline mat bsx_add_pw(const colvec& col, const rowvec& row)
{
	uword m = col.n_rows;
	uword n = row.n_cols;

	return arma::repmat(col, 1, n) + arma::repmat(row, m, 1);
}


struct sqr_fun
{
	typedef double result_type;

	double operator() (double x) { return x * x; }
};

struct sqrp_fun
{
	typedef double result_type;

	double operator() (double x, double y) { return x * x + y; }
};


struct norm_fun
{
	typedef double result_type;

	double operator() (const colvec& x) const
	{
		return arma::norm(x, 2);
	}
};

struct dotp_fun
{
	typedef double result_type;

	double operator() (const colvec& x, const colvec& y) const
	{
		return arma::dot(x, y + 2);
	}
};


/************************************************
 *
 *  Test cases
 *
 ************************************************/

TEST( ArmaBase, Reduce )
{
	const uword m = 4;
	const uword n = 6;

	mat X = arma::randu<mat>(m, n);
	mat X2 = arma::randu<mat>(m, n);
	colvec v = arma::randu<colvec>(n);
	colvec v2 = arma::randu<colvec>(n);
	double tol = 1.0e-12;

	double ma1_gt = arma::accu( X % X );
	double ma1 = map_accu(sqr_fun(), X);
	EXPECT_NEAR( ma1_gt, ma1, tol );

	double ma2_gt = arma::accu( X % X + X2 );
	double ma2 = map_accu(sqrp_fun(), X, X2);
	EXPECT_NEAR( ma2_gt, ma2, tol );

	double va1_gt = arma::sum( v % v );
	double va1 = map_sum(sqr_fun(), v);
	EXPECT_NEAR( va1_gt, va1, tol );

	double va2_gt = arma::sum( v % v + v2);
	double va2 = map_sum(sqrp_fun(), v, v2);
	EXPECT_NEAR( va2_gt, va2, tol );

	rowvec ms10_gt = arma::sum(X % X, 0);
	rowvec ms10 = map_sum(sqr_fun(), X, 0);
	ASSERT_TRUE( is_same_size(ms10, ms10_gt) );
	EXPECT_TRUE( is_approx(ms10, ms10_gt, tol) );

	colvec ms11_gt = arma::sum(X % X, 1);
	colvec ms11 = map_sum(sqr_fun(), X, 1);
	ASSERT_TRUE( is_same_size(ms11, ms11_gt) );
	EXPECT_TRUE( is_approx(ms11, ms11_gt, tol) );

	rowvec ms20_gt = arma::sum(X % X + X2, 0);
	rowvec ms20 = map_sum(sqrp_fun(), X, X2, 0);
	ASSERT_TRUE( is_same_size(ms20, ms20_gt));
	EXPECT_TRUE( is_approx(ms20, ms20_gt, tol));

	colvec ms21_gt = arma::sum(X % X + X2, 1);
	colvec ms21 = map_sum(sqrp_fun(), X, X2, 1);
	ASSERT_TRUE( is_same_size(ms21, ms21_gt));
	EXPECT_TRUE( is_approx(ms21, ms21_gt, tol));

}


TEST( ArmaBase, BsxFuns )
{
	const uword m = 4;
	const uword n = 5;

	mat X = arma::randu<mat>(m, n);
	rowvec row = arma::randu<rowvec>(n) + 0.5;
	colvec col = arma::randu<colvec>(m) + 0.5;

	double tol = 1.0e-13;

	// bsxr

	mat Ar0 = bsx_add_rows(X, row);
	mat Ar1 = bsxr_plus(X, row);

	ASSERT_TRUE( is_same_size(Ar0, Ar1) );
	EXPECT_TRUE( is_approx(Ar0, Ar1, tol) );

	mat Sr0 = bsx_sub_rows(X, row);
	mat Sr1 = bsxr_minus(X, row);

	ASSERT_TRUE( is_same_size(Sr0, Sr1) );
	EXPECT_TRUE( is_approx(Sr0, Sr1, tol) );

	mat Mr0 = bsx_mul_rows(X, row);
	mat Mr1 = bsxr_times(X, row);

	ASSERT_TRUE( is_same_size(Mr0, Mr1) );
	EXPECT_TRUE( is_approx(Mr0, Mr1, tol) );

	mat Dr0 = bsx_div_rows(X, row);
	mat Dr1 = bsxr_divides(X, row);

	ASSERT_TRUE( is_same_size(Dr0, Dr1) );
	EXPECT_TRUE( is_approx(Dr0, Dr1, tol) );

	// bsxc and bsxc_inplace

	mat Ac0 = bsx_add_cols(X, col);
	mat Ac1 = bsxc_plus(X, col);

	ASSERT_TRUE( is_same_size(Ac0, Ac1) );
	EXPECT_TRUE( is_approx(Ac0, Ac1, tol) );

	mat Sc0 = bsx_sub_cols(X, col);
	mat Sc1 = bsxc_minus(X, col);

	ASSERT_TRUE( is_same_size(Sc0, Sc1) );
	EXPECT_TRUE( is_approx(Sc0, Sc1, tol) );

	mat Mc0 = bsx_mul_cols(X, col);
	mat Mc1 = bsxc_times(X, col);

	ASSERT_TRUE( is_same_size(Mc0, Mc1) );
	EXPECT_TRUE( is_approx(Mc0, Mc1, tol) );

	mat Dc0 = bsx_div_cols(X, col);
	mat Dc1 = bsxc_divides(X, col);

	ASSERT_TRUE( is_same_size(Dc0, Dc1) );
	EXPECT_TRUE( is_approx(Dc0, Dc1, tol) );

	// pairwise

	mat Ap0 = bsx_add_pw(col, row);
	mat Ap1 = pairwise(std::plus<double>(), col, row);

	ASSERT_TRUE( is_same_size(Ap0, Ap1) );
	EXPECT_TRUE( is_approx(Ap0, Ap1, tol) );

}


TEST( ArmaBase, VecMap )
{
	const uword d = 3;
	const uword m = 4;
	const uword n = 5;

	double tol = 1.0e-12;

	mat X = arma::randu<mat>(d, m);
	mat X2 = arma::randu<mat>(d, m);
	mat Y = arma::randu<mat>(d, n);

	rowvec sw0 = arma::sqrt(arma::sum( X % X, 0 ));
	rowvec sw = colwise_vecmap(norm_fun(), X);
	ASSERT_TRUE( is_same_size(sw0, sw) );
	EXPECT_TRUE( is_approx(sw0, sw, tol) );

	rowvec cw0 = arma::sum(X % (X2 + 2), 0);
	rowvec cw = colwise_vecmap2(dotp_fun(), X, X2);
	ASSERT_TRUE( is_same_size(cw0, cw) );
	EXPECT_TRUE( is_approx(cw0, cw, tol) );

	mat P0 = X.t() * (X + 2);
	mat P = pairwise_vecmap(dotp_fun(), X, X);
	ASSERT_TRUE( is_same_size(P0, P) );
	EXPECT_TRUE( is_approx( P0, P, tol) );

	mat Q0 = X.t() * (Y + 2);
	mat Q = pairwise_vecmap(dotp_fun(), X, Y);
	ASSERT_TRUE( is_same_size(Q0, Q) );
	EXPECT_TRUE( is_approx( Q0, Q, tol) );
}




